package com.mchat.ratelimit.aspect;

import com.mchat.constant.RedisConstant;
import com.mchat.exceptions.CommonException;
import com.mchat.ratelimit.annotation.RateLimit;
import com.mchat.ratelimit.dto.BaseDTO;
import com.mchat.ratelimit.strategy.AbstractRateLimitStrategy;
import com.mchat.ratelimit.strategy.StrategyFactory;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.CodeSignature;
import org.aspectj.lang.reflect.MethodSignature;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.lang.reflect.Method;
import java.util.Arrays;


@Slf4j
@Aspect
public class RateLimitAspect {


    @Around("@annotation(com.mchat.ratelimit.annotation.RateLimit)")
    public Object aroundF(ProceedingJoinPoint pjp) {
        Method method = ((MethodSignature) pjp.getSignature()).getMethod();
        RateLimit annotation = method.getAnnotation(RateLimit.class);

        Object[] args = pjp.getArgs();
        StringBuilder prefix = new StringBuilder(32);

        String paramName = annotation.param();

        //如果没有指定限流的对象，则对象是方法的所有调用者
        StringBuilder key = new StringBuilder(64);
        //获取前缀作为唯一标识
        getPrefix(method, annotation.prefix(), prefix);

        key.append(RedisConstant.RATE_LIMIT)
                .append(':')
                .append(prefix);
        //获取arg
        getArg(pjp, args, paramName, key);
        //构建dto
        BaseDTO dto = new BaseDTO();
        dto.setCount(annotation.count());
        dto.setKey(key);
        dto.setSecond(annotation.second());
        AbstractRateLimitStrategy abstractRateLimitStrategy = StrategyFactory.get(annotation.strategyName());
        return abstractRateLimitStrategy.reachRateLimit(dto)
                .flatMapMany(isRate -> {
                    if (isRate) {
                        return Flux.error(new CommonException("请求过于频繁"));
                    }
                    try {
                        Object result = pjp.proceed(pjp.getArgs());

                        if (result instanceof Mono) {
                            return ((Mono<?>) result).flux();  // 转换为 Flux
                        } else if (result instanceof Flux) {
                            return (Flux<?>) result;
                        } else {
                            return Flux.just(result); // 普通对象包装成 Flux
                        }
                    } catch (Throwable e) {
                        return Flux.error(e);
                    }
                })
                // 兼容 Mono<T> 和 Flux<T> 的返回类型，最后再判断一下
                .as(flux -> {
                    // 如果目标方法原来是 Mono，就还原成 Mono
                    MethodSignature methodSig = (MethodSignature) pjp.getSignature();
                    Class<?> returnType = methodSig.getReturnType();
                    if (Mono.class.isAssignableFrom(returnType)) {
                        return flux.next();  // 取第一个元素还原成 Mono
                    } else {
                        return flux;        // 原来就是 Flux，直接返回
                    }
                });
    }


        private static void getPrefix (Method method, String annotationPrefix, StringBuilder prefix){
            //如果没指定前缀就用方法的全名作为前缀
            if (!annotationPrefix.equals("")) {
                prefix.append(annotationPrefix);
            } else {
                prefix.append(method.getDeclaringClass().getName())
                        .append('#')
                        .append(method.getName())
                        .append(Arrays.toString(method.getParameterTypes()));
            }
        }

        private static void getArg (ProceedingJoinPoint pjp, Object[]args, String paramName, StringBuilder key){
            if (!paramName.equals("")) {
                String[] paramNames = ((CodeSignature) pjp.getSignature()).getParameterNames();
                Object arg = null;
                for (int i = 0; i < args.length; i++) {
                    if (paramName.equals(paramNames[i])) {
                        arg = args[i];
                    }
                }
                if (arg != null)
                    key.append(':').append(arg);
            }
        }


    }
